using LinearAlgebra

function legendregauss(m)
    # Gauss Legrendre point
    β = @. .5/sqrt(1-(2*(1:m-1))^(-2.))
    T = SymTridiagonal(zeros(m), β)
    D, V = eigen(T)
    i = sortperm(D); x = D[i]
    x
end

function legendregausslobatto(m)
    # Compute the m'th order Legendre-Gauss-Lobatto quadrature points, x, and weights, w
    x = zeros(m+1); w = zeros(m+1)

    if m==1
        x[1]=-1.0; x[2]=1.0
        w[1]= 1.0; w[2]=1.0
        return x
    end

    if m==2
        x[1]=-1.0; x[2]=0.0; x[3]=1.0
        w[1]= 1/3; w[2]=4/3; w[3]=1/3
        return x
    end

    J = zeros(m-1, m-1); h1 = 2*(0:m-2).+2
    aux = broadcast(sqrt, (1:m-2).*((1:m-2).+ 2).*((1:m-2).+1).*((1:m-2).+1)./(h1[1:m-2].+1)./(h1[1:m-2].+3))
    J[1:m-2,2:m-1] = Diagonal(2.0 ./(h1[1:m-2].+2).*aux)
    J[1,1] = 0.0; J = J + J'

    D, V = eigen(J)
    i = sortperm(D); x = [-1.0; D[i]; 1.0]
    x
end

function legendre(x,m)
    # Evaluate orhonormal m'th order Legendre Polynomial at point x
    xp = x
    # Initial values P_0(x) and P_1(x)
    PL = zeros(m+1, length(xp))
    PL[1,:]  .= sqrt(1.0/2.0)
    if  m == 0
        P = PL'
        return P
    end
    PL[2,:] = sqrt(3.0/2.0)*xp
    if  m == 1
        P = PL[m+1,:]
        return P
    end

    # Forward recurrence using the symmetry of the recurrence.
    aold = sqrt(1.0/3.0)
    for i=1:m-1
        anew = 2/(2*i+2)*sqrt((i+1)*(i+1)*(i+1)*(i+1)/(2*i+1)/(2*i+3))
        PL[i+2,:] = 1/anew*(-aold*PL[i,:] + xp.*PL[i+1,:])
        aold = anew
    end
    P = PL[m+1,:]
    return P
end

function vandermonde(m,r)
    # Initialize the 1D Vandermonde Matrix, V_{ij} = phi_j(r_i)
    V = zeros(length(r),m+1)
    for j in 1:m+1
        V[:,j] = legendre(r, j-1)
    end
    V
end

function gradlegendre(r, m)
    """
    Evaluate the derivative of the
    m'th order Legendre polynomial at points r
    """
    dP = zeros(length(r),1)
    if m>0
        Ph = -m*r.*legendre(r,m) + m*sqrt((2*m+1)/(2*m-1))*legendre(r,m-1)
        dPe = r.^(m+1)*m*(m+1)/2*sqrt((2*m+1)/2)
        endp = (broadcast(abs, broadcast(abs,r).-1)) .> 10*2.2204e-16
        rh = r.*endp
        dP = broadcast(~,endp).*dPe + endp.*Ph./(1.0 .- rh.^2)
    end
    dP
end

function gradvandermonde(m,r)
    """
    Aim: Initialize the gradient of the Vandermonde matrix of order m at (r)
    """
    Vr = zeros(length(r),m+1)
    for i=0:m
       Vr[:,i+1] = gradlegendre(r,i)
    end
    Vr
end

function dmatrix(m,r,V)
    """
    Purpose : Initialize the (r) differentiation matrices,
    evaluated at (r) at order m
    """
    Vr = gradvandermonde(m, r); D = Vr/V
    D
end

function geomatry(m,N)
    """
    geomatry informations
    """
    #Initialize extraction vector
    VtoE = zeros(Int64, 2, N)
    for j = 1:N
        VtoE[1, j] = (j - 1) * (m + 1) + 1
        VtoE[2, j] = j * (m + 1)
    end

    vmapM = VtoE[:]
    vmapP = zeros(Int64, 2, N); vmapP[1,1] = VtoE[1,1]; vmapP[2,N] = VtoE[2,N]
    vmapP[2,1:N-1]= VtoE[1,2:N]; vmapP[1,2:N]=VtoE[2,1:N-1]
    vmapP = vmapP[:]
    vmapI = vmapM[1]; vmapO = vmapM[end]
    vmapM, vmapP, vmapI, vmapO
end

function matrics(m,r,x)
    # Construct Operators at [-1,1]
    V = vandermonde(m, r); D = dmatrix(m, r, V)
    M = inv(V * V'); S = M * D

    xr  = D*x; J = xr; rx = 1 ./ J; Fscale = 1 ./ (J[[1, m+1],:])

    Emat = zeros(m+1, 2); Emat[1,1] = 1.0; Emat[m+1,2] = 1.0
    LIFT = V*(V'*Emat)
    D, rx, LIFT, Fscale
end

function rhs(time,ϵ,xL,xR,u,h,m,N,D,vmapM,vmapP,vmapI,vmapO,rx, LIFT, Fscale)

    """
    Calculate the right hand side by weak form
    """
    # define field differences at faces
    du = zeros(2, N); du = reshape(u[:][vmapM] - u[:][vmapP], 2, N)
    # impose boundary conditions
    mapI = 1; mapO = 2*N
    uin  = -tanh((xL+0.5-time)/(2*ϵ))+1.0; du[mapI] = 2.0*(u[vmapI] - uin)
    uout = -tanh((xR+0.5-time)/(2*ϵ))+1.0; du[mapO] = 2.0*(u[vmapO] - uout)

    # Compute q and jumps
    nx = zeros(2, N); nx[1,:] .= -1; nx[2,:] .= 1
    q = sqrt(ϵ)*(rx .* (D*u)) - LIFT*(Fscale .* (nx .* du ./ 2.0))
    dq = zeros(2, N); dq = reshape(q[:][vmapM] - q[:][vmapP], 2, N) ./ 2

    # impose boundary condition - Dirichlet conditions
    dq[mapI] = 0.0; dq[mapO] = 0.0

    # evaluate nonlinear flux
    du2 = zeros(2, N); du2 = reshape((u[:][vmapM]).^2 - (u[:][vmapP]).^2, 2, N) ./ 2

    # impose boundary condition
    du2[mapI] = (u[vmapI]^2 - uin^2); du2[mapO] =  (u[vmapO].^2 - uout.^2)

    # compute flux
    maxvel = maximum(maximum(abs.(u), dims=2))

    # penalty scaling
    tau = 0

    # flux term
    flux = nx.*(du2 ./ 2.0 - sqrt(ϵ)*dq) - maxvel/2.0 .* du - sqrt(ϵ)*tau .* du

    # local derivatives of field
    dfdr = D*(u.^2 ./ 2.0 - sqrt(ϵ)*q)

    # compute right hand sides of the semi-discrete PDE
    rhsu = -(rx .* dfdr - LIFT*(Fscale .* flux))

    rhsu
end
